import os, torch
from copy import deepcopy
import numpy as np
from torch.utils.data import Dataset
from .data_utils import Helper
torch.set_default_dtype(torch.float32)

class Dataset2D(Dataset):
    def __init__(self, path, eof_name, transforms=None, loader_fnc=Helper.loadTensor):
        super().__init__()
        self.path = path
        self.loader = loader_fnc # Load tensor with type of Float32

        print('Loading Dataset')
        self._volume = self.loader(path + '/vol_' + eof_name + '.pth')
        self._mask = self.loader(path + '/mask_' + eof_name + '.pth')
        print('Loading completed !')

    def __getitem__(self, i):
        return self._volume[i], self._mask[i]

    def __len__(self):
        return self._volume.shape[0]

class Dataset3D(Dataset):
    def __init__(self, path, transforms=None, loader_fnc=Helper.loadTensor):
        super().__init__()
        self.path = path
        self.loader = loader_fnc # Load tensor with type of Float32
        self.len = int(self.loader(path + '/info.pth').item())
        
    def __getitem__(self, i):
        # load data
        self._vol = self.loader(self.path + '/volumes/vol_' + str(i) + '.pth')
        self._mask = self.loader(self.path + '/masks/vol_' + str(i) + '.pth')

        return self._vol, self._mask

    def __len__(self):
        return self.len
